import torch.nn as nn
import torch.nn.functional as F

from ...losses import BuildLoss
from ..base import BaseSegmentor
from ...backbones import BuildNormalization, BuildActivation
from .head import DDRHead, DDRAUCHead

'''DDRNet'''
class DDRNet(BaseSegmentor):
    def __init__(self, cfg, mode):
        super(DDRNet, self).__init__(cfg, mode)
        align_corners, norm_cfg, act_cfg, head_cfg = self.align_corners, self.norm_cfg, self.act_cfg, cfg['head']
        self.decoder = DDRHead(head_cfg['in_channels'], head_cfg['feats_channels'], cfg['num_classes'], head_cfg['dropout_ratio'], head_cfg['align_corners'])
        # build auxiliary decoder
        self.setauxiliarydecoder(cfg['auxiliary'])
    '''forward'''
    def forward(self, x, targets=None):
        img_size = x.size(2), x.size(3)
        # feed to backbone network
        backbone_outputs = self.backbone_net(x)
        # feed to decoder
        predictions = self.decoder(backbone_outputs[self.cfg['head']['in_index']])
        # forward according to the mode
        if self.mode == 'TRAIN':
            loss, losses_log_dict = self.customizepredsandlosses(
                predictions=predictions, targets=targets, backbone_outputs=backbone_outputs, losses_cfg=self.cfg['losses'], img_size=img_size,auxiliary_cfg=self.cfg['auxiliary']
            )
            return loss, losses_log_dict
        return predictions

    def setauxiliarydecoder(self, auxiliary_cfg):
        norm_cfg, act_cfg, num_classes = self.norm_cfg.copy(), self.act_cfg.copy(), self.cfg['num_classes']
        if auxiliary_cfg is None: return
        if isinstance(auxiliary_cfg, dict):
            auxiliary_cfg = [auxiliary_cfg]
        self.auxiliary_decoder = nn.ModuleList()
        for aux_cfg in auxiliary_cfg:
            if aux_cfg['type'] == "DDRAUCHead":
                dec = DDRAUCHead(aux_cfg['in_channels'], aux_cfg['feats_channels'], self.cfg['num_classes'], aux_cfg['dropout_ratio'], aux_cfg['align_corners'])
            self.auxiliary_decoder.append(dec)
        # if len(self.auxiliary_decoder) == 1:
        #     self.auxiliary_decoder = self.auxiliary_decoder[0]

    def customizepredsandlosses(self, predictions, targets, backbone_outputs, losses_cfg, img_size, auxiliary_cfg,auto_calc_loss=True, map_preds_to_tgts_dict=None):
        predictions = F.interpolate(predictions, size=img_size, mode='bilinear', align_corners=self.align_corners)
        outputs_dict = {'loss_cls': predictions}
        for idx, (aux_cfg, dec) in enumerate(zip(auxiliary_cfg,self.auxiliary_decoder)):
            backbone_output = backbone_outputs[aux_cfg["in_index"]]
            predictions_aux = dec(backbone_output)
            if aux_cfg['type'] == "DDRAUCHead":
                predictions_aux = F.interpolate(predictions_aux, size=img_size, mode='bilinear', align_corners=self.align_corners)
            outputs_dict[aux_cfg['type']] = predictions_aux
        return self.calculatelosses(predictions=outputs_dict, targets=targets, losses_cfg=losses_cfg, map_preds_to_tgts_dict=map_preds_to_tgts_dict)
